import time
import socket

from typing           import Dict, Any, Optional

from pyverbs.addr     import GID                 # type: ignore
from pyverbs.mr       import MR                  # type: ignore
from pyverbs.cq       import WC                  # type: ignore
from pyverbs.cm_enums import RAI_PASSIVE         # type: ignore
from pyverbs.cm_enums import RDMA_PS_TCP         # type: ignore
from pyverbs.cmid     import CMID, AddrInfo      # type: ignore
from pyverbs.qp       import QPInitAttr, QPCap   # type: ignore


# 为容错保留的解控缓存空间
RESERVED_LEN      = 20
HANDSHAKE_WORDS   = "Aloha!"   # 握手通信的消息内容


# 自定义异常
class CommError(Exception):
    def __init__(self, message):
        self.message = message
        super(CommError, self).__init__(message)

    def __str__(self):
        return self.message
    



class ConnManager:
    def __init__(self):
        pass

    @staticmethod
    def prepare_send_msg(**kwargs) -> str:
        if not bool(kwargs):
            send_msg = HANDSHAKE_WORDS
        else:
            print("-- 本地信息")
            send_msg = ""
            for key, value in kwargs.items():
                print(key, ":", value)
                send_msg = send_msg + key + ':'   \
                        + type(value).__name__    \
                        + ':' + str(value) + ','

            send_msg = send_msg[:-1]
            print('-' * 80)

        return send_msg

    @staticmethod
    def parse_recv_msg(only_handshake: bool, 
                recv_msg: str) -> Dict[Any, Any]:
        if only_handshake:
            try:
                if recv_msg != HANDSHAKE_WORDS:
                    raise CommError("Failed to handshake with remote "
                            "peer by " + ConnManager.__class__.__name__)
            except CommError as e:
                print(e)
        else:
            key_value = {}
            print("-- 远端信息")
            for item in recv_msg.split(','):
                key, value_type, value = item.split(':', 2)
                key_value[key.strip()] = eval(value_type.strip())(value.strip())
                print(key, ":", value)
            print('-' * 80)

            return key_value
        



class SocketCM(ConnManager):
    def __init__(self, port: int, ip: Optional[str]=None) -> None:
        """使用套接字进行建连"""
        super(ConnManager, self).__init__()

        self.socket = socket.socket()
        # 设置套接字重用
        self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.is_server = False

        # 未指定IP表示为服务端(RDMA网卡同时支持以太网和RDMA协议)
        if ip is None:  
            self.is_server = True
            # "0.0.0.0" 表示可接收任意客户端IP的连接
            self.socket.bind(('0.0.0.0', port))  
            self.socket.listen()

            # 阻塞等待一个请求进来
            # 将套接字设为客户端套接字(这里只允许单个连接)
            self.client, self.addr = self.socket.accept()
            self.socket.close()
            self.socket = self.client
        else:   # 客户端，连接到服务端
            self.socket.connect((ip, port))


    def handshake(self, **kwargs) -> Dict[Any, Any]:
        """通过套接字握手进行元数据交换，比如RDMA的MR信息、秘钥信息等
        For example:
            1. socket.handshake()
                This exchanges nothing, which is always used as a synchronization mechanism.
            # QPN: QP的编号, QP Number
            # PSN: 数据包序列号,Packet Sequence Number
            2. socket.handshake(qpn = xx, psn = yy)
                This exchanges qpn and psn with remote peer, this can be used before INIT->RTR.
            3. socket.handshake(str = "hello bro")
                This exchanges a message between server and client.

        :return: 从Peer端接收到的元信息
        """
        just_handshake = not bool(kwargs)               # 关键字参数为空

        send_msg = self.prepare_send_msg(**kwargs)      # 准备握手发送的消息
        print(f"send_msg = {send_msg}")

        self.socket.send(send_msg.encode('utf-8'))      # 编码为bytes类型再发送

        # 接收响应信息
        recv_msg = self.socket.recv(len(send_msg) + RESERVED_LEN).decode('utf-8')
        print(f"recv_msg = {recv_msg}")

        time.sleep(0.5)

        return self.parse_recv_msg(just_handshake, recv_msg)



    def close(self):
        if not self.is_server:
            self.socket.close()



    def __del__(self):
        self.socket.close()




class RDMACM(ConnManager):
    def __init__(self, port: int, ip: Optional[str] = None) -> None:
        """使用RDMA自带连接管理进行建连"""
        super(ConnManager, self).__init__()

        # QPCap用于定于QP的能力， 支持的参数有：
        #   - max_send_wr     最大发送WR(Work Request, 工作请求，驱动层面体现为WQE)
        #   - max_recv_wr     最大接收WR
        #   - max_send_sge    每个发送WR中允许的最大发送SGE
        #   - max_recv_sge    每个发送WR中允许的最大接收SGE
        #   - max_inline_data 内联数据的最大大小（字节），若支持内联发送，数据可直接嵌入WR无需DMA。
        #! Ref：https://github.com/linux-rdma/rdma-core/blob/master/pyverbs/qp.pyx#L34
        # TODO(yangxianpku: 2025.08.04): 考虑调优
        cap            = QPCap(max_recv_wr=1)


        #! QPInitAttr用于定义QP的初始属性， 参数包括:
        #   - qp_type     QP类型, 支持IBV_QPT_RC(可靠连接)， IBV_QPT_UC(不可靠连接)， I
        #                   BV_QPT_UD(不可靠数据包)         
        #   - qp_context  QPC, 即QP上下文
        #   - scq         发送完成队列
        #   - rcq         接收完成队列
        #   - srq         共享接收队列
        #   - cap         QPCap对象
        #   - sq_sig_all  是否所有WR都需要生成CQE, 默认True
        #! Ref：https://github.com/linux-rdma/rdma-core/blob/master/pyverbs/qp.pyx#L100
        qp_init_attr   = QPInitAttr(cap=cap)
        self.is_server = False


        # TODO: Create a QP after connection instead of creating one each time user want to handshake

        if ip is None:   # 服务端
            self.is_server = True

            # AddrInfo创建地址信息
            #  - src                 服务端绑定的IP地址 
            #  - dst                 客端连接的目的地址 
            #  - src_service         源地址的服务名或端口号 
            #  - dst_service         目标地址的服务名或端口号
            #  - port_space=0,       RDMA端口空间， RDMA_PS_TCP或RDMA_PS_UDP
            #  - flags=0             - RAI_PASSIVE     被动模式, 服务端监听连接请求 
            #                        - RAI_NUMERICHOST 使用IP地址（如 "192.168.1.1"），禁止主机名解析
            #                        - RAI_NOROUTE     不检查路由可达性，即跳过对目标地址的路由表校
            #                        - RAI_FAMILY      指定 地址族（Address Family），强制使用IPv4或IPv6
            #                        - 通过'|'传递多个值
            #! Ref：https://github.com/linux-rdma/rdma-core/blob/master/pyverbs/cmid.pyx#L167
            cai = AddrInfo(src         = '0.0.0.0', 
                           src_service = str(port), 
                           port_space  = RDMA_PS_TCP, 
                           flags       = RAI_PASSIVE
                        )
            
            #! 创建通信管理器标识符(Communication Manager Identifier, CMID)
            # 主要作用: 
            #   - 连接管理:      封装QP, PD等资源，管理连接状态机
            #   - 地址解析与路由: 通过rdma_resolve_addr解析目标地址，rdma_resolve_route确定路由路径
            #   - 事件处理:      监听连接事件（如连接请求、建立完成、断开），异步通知应用程序
            #   - 多协议支持:    支持RC（可靠连接）、UC（不可靠连接）、UD（不可靠数据报）等传输类型
            #
            # [Note]: 可以理解为RDMA版本的套接字
            #! Ref：https://github.com/linux-rdma/rdma-core/blob/master/pyverbs/cmid.pyx#L323
            self.cmid = CMID(creator     = cai, 
                            qp_init_attr = qp_init_attr
                        )
            self.cmid.listen()

            client_cmid = self.cmid.get_request()
            client_cmid.accept()
            self.cmid.close()
            self.cmid   = client_cmid
        else:          # 客户端
            cai       = AddrInfo(src        = '0.0.0.0', 
                                dst         = ip, 
                                dst_service = str(port), 
                                port_space  = RDMA_PS_TCP, 
                                flags       = 0,
                            )
            self.cmid = CMID(creator=cai, qp_init_attr=qp_init_attr)
            self.cmid.connect()

    def handshake(self, **kwargs):
        """通过CM握手进行元数据交换，比如RDMA的MR信息、秘钥信息等
        For example:
            1. socket.handshake()
                This exchanges nothing, which is always used as a synchronization mechanism.
            2. socket.handshake(qpn = xx, psn = yy)
                This exchanges qpn and psn with remote peer, this can be used before INIT->RTR.
            3. socket.handshake(str = "hello bro")
                This exchanges a message between server and client.
        :return: 从Peer端接收到的元信息
        """
        just_handshake = not bool(kwargs)

        send_msg = self.prepare_send_msg(**kwargs)
        size     = len(send_msg)

        # print(f"send_msg = {send_msg}")

        #! Recv Message
        # 1. reg_msgs: 注册内存区域用于接收消息
        recv_mr: MR = self.cmid.reg_msgs(size + RESERVED_LEN)
        # 2. 下发一个接收请求, 告诉RDMA HCA硬件将接收到的数据存放至recv_mr
        # recv_mr: MR   接收数据的MR
        # length: int   MR区域的长度
        #! Ref： post_send   https://github.com/linux-rdma/rdma-core/blob/master/pyverbs/cmid.pyx#L693
        self.cmid.post_recv(recv_mr, size)

        #! Send Message
        # 1. reg_msgs: 注册内存区域用于发送消息
        send_mr: MR = self.cmid.reg_msgs(size)

        # 2. 将用户数据发送至MR缓冲区(memcpy)
        send_mr.write(send_msg.encode('utf-8'), size)

        # 3. 下发一个发送请求, 告诉RDMA HCA硬件将接收到的数据存放至recv_mr
        # send_mr: MR   发送数据的MR
        # flags:  enum  发送标签, 支持:
        #                   - IBV_SEND_FENCE
        #                   - IBV_SEND_SIGNALED
        #                   - IBV_SEND_SOLICITED
        #                   - IBV_SEND_INLINE
        #                   - IBV_SEND_IP_CSUM
        # length: int   MR区域的长度
        #! Ref： post_send       https://github.com/linux-rdma/rdma-core/blob/master/pyverbs/cmid.pyx#L707
        #! Ref： ibv_send_flags  https://github.com/linux-rdma/rdma-core/blob/master/pyverbs/libibverbs_enums.pxd#L157
        self.cmid.post_send(send_mr, 0, size)

        #! Ref： https://github.com/linux-rdma/rdma-core/blob/master/pyverbs/cmid.pyx#L771
        recv_wc: WC = self.cmid.get_recv_comp()  # Poll一个接收CQ
        send_wc: WC = self.cmid.get_send_comp()  # Poll一个发送CQ

        # 从接收缓冲区中读取数据
        recv_msg = recv_mr.read(recv_wc.byte_len, 0).decode('utf-8')

        # print(f"recv_msg = {recv_msg}")


        return self.parse_recv_msg(just_handshake, recv_msg)

    def close(self):
        if not self.is_server:
            self.cmid.close()

    def __del__(self):
        self.cmid.close()